import cv2
import numpy as np
from conf import *
from utils import *

color_list = ['green', 'blue', 'yellow']

tmp_color = 'yellow'
if tmp_color in color_list:
    color_lower, color_upper = set_color(tmp_color)
elif tmp_color == 'red':
    color_lower1, color_upper1, color_lower2, color_upper2 = set_color(tmp_color)
else:
    raise IndexError


def main():
    cap = set_camera(width, height)
    while True:
        _, frame = cap.read()
        frame = cv2.GaussianBlur(frame, (3, 3), 0)
        hsv = cv2.cvtColor(frame, cv2.COLOR_BGR2HSV)
        if tmp_color in color_list:
            mask = cv2.inRange(hsv, color_lower, color_upper)
        elif tmp_color == 'red':
            mask1 = cv2.inRange(hsv, color_lower1, color_upper1)
            mask2 = cv2.inRange(hsv, color_lower2, color_upper2)
            mask = cv2.bitwise_or(mask1, mask2)
        else:
            raise IndexError
        mask = cv2.erode(mask, np.ones((5, 5), np.uint8), iterations=2)
        # mask = cv2.dilate(mask, np.ones((5, 5), np.uint8), iterations=2)
        mask = cv2.GaussianBlur(mask, (3, 3), 0)
        cnts = cv2.findContours(mask.copy(), cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)[-2]

        ls = []
        if len(cnts) > 0:  # 通过边缘检测来确定所识别物体的位置信息得到相对坐标
            for i in range(len(cnts)):
                cnt = cnts[i]
                cnt = cv2.approxPolyDP(cnt, 0.05 * cv2.arcLength(cnt, True), True)
                if cnt.shape[0] == 4:
                    ls.append(cnt)

        if len(ls):
            for cnt in ls:
                if is_square(cnt):
                    rect = cv2.minAreaRect(cnt)
                    box = cv2.boxPoints(rect)
                    # box = cnt.squeeze()
                    cv2.drawContours(frame, [np.int0(box)], -1, (0, 255, 255), 2)

        cv2.imshow('frame', frame)  # 将具体的测试效果显示出来
        cv2.imshow('mask', mask)
        if cv2.waitKey(1) & 0xFF == 27:  # 如果按了ESC就退出 当然也可以自己设置
            break

    cap.release()
    cv2.destroyAllWindows()  # 后面两句是常规操作,每次使用摄像头都需要这样设置一波


if __name__ == '__main__':
    main()
